import os

import cv2
import numpy as np
import torch

from models.experimental import attempt_load
from utils.augmentations import letterbox
from utils.general import non_max_suppression, scale_coords, xyxy2xywh
from utils.plots import Annotator

device = 'cuda' if torch.cuda.is_available() else 'cpu'
half = device != 'cpu'
# 获取权重
weights = r'E:\2021homework\cv_masks\yolo\runs\train\exp9\weights\best.pt'
# weights = r'work/weights/best.pt'
imgsz = 640
conf_thres = 0.4
iou_thres = 0.2


def load_model():
    model = attempt_load(weights, map_location=device)  # load FP32 model

    if half:
        model.half()  # to FP16

    if device != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
    return model


def create_video_writer(input_video, output_path):
    """Creates a OpenCV Video Writer object using the mp4c codec and
    input video stats (frame width, height, fps) for tracking
    visualization.

    Args:
      input_video: video read into program using opencv methods

    Returns:
      out: cv2 videowriter object to which individual frames can be
           written
    """
    # Grab some video stats for videowriter object.
    original_width = input_video.get(cv2.CAP_PROP_FRAME_WIDTH)
    original_height = input_video.get(cv2.CAP_PROP_FRAME_HEIGHT)
    fps = input_video.get(cv2.CAP_PROP_FPS)
    file_name = 'output.mp4'
    output_path = output_path + os.sep + file_name
    # Make an output directory if necessary.

    # Initiate video writer object by defining the codec and initiating
    # the VideoWriter object.
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path,
                          fourcc,
                          fps,
                          (int(original_width), int(original_height)),
                          isColor=True)

    return out


if __name__ == '__main__':
    # 获取视频流以及相关初始化
    cap = cv2.VideoCapture(0)
    model = load_model()
    stride = int(model.stride.max())
    names = model.module.names if hasattr(model, 'module') else model.names
    frame_num = 0
    output_path = r'E:\2021homework\cv_masks\yolo\output'
    out = create_video_writer(cap, output_path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not (ret):
            continue
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        frame_fps = int(cap.get(cv2.CAP_PROP_FPS))
        s_width, s_height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        # cv2.imshow("cap", frame)

        img0 = frame
        src = img0
        # pil_img = Image.open(frame)

        # img0 = cv2.cvtColor(img0, cv2.COLOR_RGB2BGR)
        img = letterbox(img0, imgsz, stride=stride)[0]

        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)

        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()
        if len(img.shape) == 3:
            img = img[None]
        img = img / 255.
        pred = model(img, augment=False, visualize=False)[0]
        pred = non_max_suppression(pred, conf_thres, iou_thres, agnostic=False, max_det=300)
        # print(pred)
        aims = []
        # print(pred)
        for i, det in enumerate(pred):
            # print(det)
            s = ''
            # 这里使用的是letterbox()改变后的img尺寸
            s += '%gx%g ' % img.shape[2:]
            # torch.tensor()使用的是原图的尺寸
            gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]
            annotator = Annotator(src, line_width=3, pil=not ascii)
            if len(det):
                # Rescale boxes from img_size to im0 size
                # 第一次将下列代码的img0写成了img QAQ
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
                """
                0:no_wear, 1:wear_right, 2:wear_wrong

                """
                # save bbox
                for *xyxy, conf, cls in reversed(det):  # Write to file
                    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                    line = (cls, *xywh)  # label format
                    aim = ('%g ' * len(line)).rstrip() % line
                    aim = aim.split(' ')
                    # print(aim)
                    aims.append(aim)

        # s_width, s_height = pil_img.size[0], pil_img.size[1]

        if len(aims):
            # print(aims)
            for i, det in enumerate(aims):
                tag, x_center, y_center, width, height = det
                # print(x_center, y_center)
                """x_center, width = s_width * float(x_center), s_width * float(width)
                y_center, height = s_height * float(y_center), s_height * float(height)
                # print(y_center, height)
                top_left = (int(x_center - width / 2.), int(y_center - height / 2.))
                bottle_right = (int(x_center + width / 2.), int(y_center + height / 2.))"""
                classes = ['no_wear', 'wear_right', 'wear_wrong']

                # print(top_left, bottle_right)

                """b0 = (float(x_center) + 1) * s_width - s_width * float(width) / 2.
                b1 = (float(y_center) + 1) * s_height - s_height * float(height) / 2.
                b2 = (float(x_center) + 1) * s_width + s_width * float(width) / 2.
                b3 = (float(y_center) + 1) * s_height + s_height * float(height) / 2.
                """
                b0 = float(x_center) * s_width + 1 - s_width * float(width) / 2.
                b1 = float(y_center) * s_height + 1 - s_height * float(height) / 2.
                b2 = float(x_center) * s_width + 1 + s_width * float(width) / 2.
                b3 = float(y_center) * s_height + 1 + s_height * float(height) / 2.
                top_left = (int(b0), int(b1))
                bottle_right = (int(b2), int(b3))
                # print(classes[int(tag)], top_left, bottle_right)

                # 框取人脸
                cv2.rectangle(frame, top_left, bottle_right, (255, 0, 0), thickness=2)
                cv2.putText(frame, classes[int(tag)], top_left, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                cv2.putText(frame, str(frame_fps), (0, 0), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        cv2.imshow("cap2", frame)
        out.write(frame)
        print('save', frame_num)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            cap.release()
            cv2.destroyAllWindows()
            out.release()
            break
        frame_num += 1

        # 后续在人脸方框附近添加戴口罩情况
